open Utils.Symbols
open Utils.Pervasive
open Syntax
open Binding
open Core
open PrintExpr

exception TreeFormationError 

type tree_expr =
	| Type of tree_type 
	| Term of tree_term 
 
 and tree_expr_pair = tree_expr * tree_expr
	
 and tree_type = 
	| T_TyArr of tree_expr_pair 
	| T_TyVar of type_var
	| T_TyAll of tree_expr_pair
  
  and type_var = 
	| T_FreeType of string
	| T_BoundedType of string
  
  and tree_term = inner_term * tree_type
  
  and inner_term = 
	| T_TmApp of tree_expr_pair
	| T_TmVar of term_var
	| T_TmAbs of tree_expr_pair
	| T_TmTAbs of tree_expr_pair
  
  and term_var =
	| T_FreeTerm of string
	| T_BoundedTerm of string		 
	| T_TmString of string
	| T_TmInt of int
	
	
	
	

	
let rec type_of_tree tree_ty = 
	match tree_ty with
	| T_TyAll ( Type (T_TyVar (T_BoundedType tyX)),  Type tree_ty ) -> 
		fun ctx -> 
			let ctx1 =  add_tvar_name ctx tyX in 
			(try TyAll ((type_of_tree tree_ty) ctx1) 
			with | Failure x -> prf x; failwith "\ntype_of_tree tree_ty")
			
			
	| T_TyAll _ -> prf "gh"; raise TreeFormationError
	| _ -> type_of_ArrowType_tree tree_ty

and type_of_ArrowType_tree tree_ty = 
	let is_aType = function
		| T_TyVar _ -> true
		|_-> false in
	
	if is_aType tree_ty then 
		type_of_AType_tree tree_ty
	else match tree_ty with 
	| T_TyArr (Type tree1, Type tree2) ->
		fun ctx -> 
			TyArr ((type_of_tree tree1) ctx, (type_of_tree tree2) ctx)
	| _ -> prf "ghsgh"; raise TreeFormationError


and type_of_AType_tree tree_ty = match tree_ty with
	| T_TyVar (T_FreeType str) when str = "String" -> 
		fun ctx -> TyString
	| T_TyVar (T_FreeType str) when str = s_TyBool -> 
		fun ctx -> TyBool
	| T_TyVar (T_FreeType str)  when str = s_TyFloat ->
		fun ctx -> TyFloat
	| T_TyVar (T_FreeType str) when str = "Nat" ->
		fun ctx -> TyNat 
	| T_TyVar (T_FreeType str) when str =  s_TyInt -> 
		fun ctx -> TyInt 
	| T_TyVar (T_FreeType str) when str =  s_TyUnit -> 
		fun ctx -> TyUnit 
	| T_TyVar (T_FreeType str) -> 
		fun ctx -> if is_name_bound ctx str 
		then TyVar (index_of_name ctx str)
		else raise TreeFormationError
	| T_TyVar (T_BoundedType  str) ->
		fun ctx -> if is_name_bound ctx str 
		then TyVar (index_of_name ctx str)
		else failwith ("In type_of_AType_tree, " ^str^" is unbound ")
	| _ -> prf "No case"; raise TreeFormationError	
	
	
	
	
	
	
let rec tree_of_type ctx ty = 
	match ty with
	| TyAll ty ->
		let (ctx1, tyX) = fresh_tvar ctx  in
		T_TyAll ( Type (T_TyVar (T_BoundedType tyX)),  Type (tree_of_type ctx1 ty)) 
	| _ -> tree_of_ArrowType ctx ty 

and tree_of_ArrowType ctx ty = 
	match ty with 
    | TyArr (ty1, ty2) ->
		T_TyArr (Type (tree_of_ArrowType ctx ty1) , Type (tree_of_ArrowType ctx ty2)) 
   | _ -> tree_of_AType ctx ty 
   
and tree_of_AType ctx ty  = 
	match ty with
	| TyVar x -> begin
			match get_binding ctx x with
			| BoundedTyVarBind -> T_TyVar (T_BoundedType (name_of_index ctx x))
			| _ -> T_TyVar (T_FreeType (name_of_index ctx x))
			end
    | TyString -> T_TyVar (T_FreeType s_TyString)
	| TyBool -> T_TyVar ( T_FreeType s_TyBool)
	| TyFloat ->T_TyVar ( T_FreeType s_TyFloat) 
	| TyNat -> T_TyVar ( T_FreeType "Nat")
	| TyInt -> T_TyVar (T_FreeType s_TyInt)
	| TyUnit -> T_TyVar (T_FreeType s_TyUnit)
	| _ -> tree_of_type ctx ty 		
	

let rec tree_of_term ctx ter = 
	
	let inner = match ter with
    | TmAbs (ty, te) ->
		 let (ctx1, x') = fresh_fvar ctx ty in
		 let var_node = Term (T_TmVar (T_BoundedTerm x'), tree_of_type ctx ty) in
         T_TmAbs  (var_node, tree_of_term ctx1 te)
    | TmTAbs ter ->
			let (ctx1, x) = fresh_tvar ctx in
			let var_node = Type(T_TyVar(T_BoundedType x)) in
			T_TmTAbs(var_node, tree_of_term ctx1 ter)
	| TmApp (ter1, ter2) ->
		T_TmApp (tree_of_term ctx ter1, tree_of_term ctx ter2) 
	
	| TmTApp (ter, ty) ->
		T_TmApp (tree_of_term ctx ter, Type(tree_of_type ctx ty)) 
	
	| TmVar x -> 
			begin match get_binding ctx x with
			| BoundedTeVarBind _  -> T_TmVar (T_BoundedTerm (name_of_index ctx x))
			| TeVarBind _ 
			| TmAbbBind _ ->T_TmVar (T_FreeTerm (name_of_index ctx x))
			| _ -> failwith 
					("tree_of_term::inner::TmVar, no handler for " ^ 
					string_of_term ctx ter Alpha)
			end
	| TmString str -> T_TmVar (T_TmString str)
	| TmInt num -> T_TmVar (T_TmInt num)
	| _ -> failwith "tree_of_term::inner no handler"

	in Term (inner, tree_of_type ctx (typeof ctx ter))
	
let tree_of_expr ctx expr = 
	match expr with
	| Te (ter, _) -> (try tree_of_term ctx ter with _->failwith "Te case")
	| Ty (ty, _) -> (try Type(tree_of_type ctx ty) with _->failwith "Ty case")
	
	
		
let rec term_of_tree (inner_ter, tree_ty) =
	
	match inner_ter with
	| T_TmVar tm_var -> 
		(fun ctx -> 
			match tm_var with
			| T_FreeTerm str -> TmVar (index_of_name ctx str)
			| T_BoundedTerm str -> 
				(try let ind = index_of_name ctx str in TmVar ind
				with x -> prf ("term_of_tree::T_BoundedTerm str case " ^ str ^ " is unbound\n"); raise x) 
			| T_TmString str -> TmString str
			| T_TmInt num -> TmInt num)
			
	| T_TmAbs (Term tree_ter1, Term tree_ter2) ->
		begin match tree_ter1 with
		| T_TmVar (T_BoundedTerm abs_var), tree_ty ->
				(fun ctx -> 
					try (let ty = (try type_of_tree tree_ty ctx with | x -> prf "\nError in type of tree"; raise x) in 
					let ctx1 = add_fvar_name ctx abs_var ty in
					TmAbs (ty, term_of_tree tree_ter2 ctx1) ) 
					with x -> prf "term_of_tree::T_TmAbs"; raise x) 
		| _ -> failwith "term_of_tree->T_TmVar case" end
	
	
	| T_TmTAbs (Type (T_TyVar t_var), Term tree_ter) ->
		begin match t_var with 
		| T_BoundedType name ->
			fun ctx -> 
				let ctx1 =  add_tvar_name ctx name  in
				TmTAbs (term_of_tree tree_ter ctx1) 
		| _ -> failwith "term_of_tree->T_TmTAbs case"	end
	
	| T_TmApp (Term tree_ter1, Term tree_ter2) -> 
		fun ctx -> TmApp 
			(term_of_tree tree_ter1 ctx, term_of_tree tree_ter2 ctx)
	| T_TmApp (Term tree_ter, Type tree_ty) ->
		fun ctx -> TmTApp 
			(term_of_tree tree_ter ctx, type_of_tree tree_ty ctx)
			
	|_ ->failwith "term_of_tree"
	
let expr_of_tree tree ctx = 
	match tree with
	| Term tree_ter -> 
		(try Te (term_of_tree tree_ter ctx, None) with 
			| Failure x -> failwith "Problem in term_of_tree") 
	| Type tree_ty -> Ty(type_of_tree tree_ty ctx, None)
	
	
let  free_term_vars ctx ter = 
	
	let rec inner lst tree_ter = 
	match tree_ter with 
	| Term(T_TmVar tm_var, _) -> 
		union (fun x y -> x = y)
		(match tm_var with
			| T_FreeTerm str -> [TmVar (index_of_name ctx str)]
			| T_TmString str -> [TmString str]
			| T_TmInt num -> [TmInt num]
			| _-> []
			) lst 
	| Term(T_TmAbs (_, tree_ter2) , _) ->
		inner lst tree_ter2 
	
	| Term(T_TmTAbs (t_var, tree_ter),_)->
		inner lst tree_ter 
	
	| Term(T_TmApp (Term tree_ter1, Term tree_ter2), _) -> 
		union (fun x y -> x=y) 
		(inner lst (Term tree_ter1))
		(inner lst (Term tree_ter2))
	
	| Term(T_TmApp (Term tree_ter, _), _)  ->
		inner lst (Term tree_ter)
	
	| _->  raise TreeFormationError
	
	in inner [] (tree_of_term ctx ter)
	



let rec tree_eq tree1 tree2 = 
	
	let p_tree_eq (a1, b1) (a2, b2) =
		tree_eq a1 a2 && tree_eq b1 b2 in
	
	let ty_eq tree_ty1 tree_ty2 = 
		match tree_ty1, tree_ty2 with 
		| T_TyAll (var1, body1),  T_TyAll (var2, body2) ->
			tree_eq (replace_all body1 var1 var2) body2
		| T_TyArr pair1, T_TyArr pair2 -> 
			p_tree_eq pair1 pair2
		| _-> false in
	
	
	let te_eq (te1, ty1) (te2, ty2) =
		ty_eq ty1 ty2 && (match te1, te2 with
		|T_TmAbs (var1, body1), T_TmAbs (var2, body2)  ->
			 tree_eq (replace_all body2 var2 var1) body1
		|T_TmTAbs (var1, body1), T_TmTAbs (var2, body2) -> 	
			tree_eq (replace_all body2 var2 var1) body1
		|T_TmApp(left1, right1), T_TmApp(left2, right2)->
			tree_eq left1 left2 && tree_eq right1 right2 
		|T_TmVar var1, T_TmVar var2 ->
			var1 = var2 
		|_ -> false) in
	
	
	if tree1 = tree2 then true else(
		match tree1, tree2 with
		| Term te1, Term te2 ->
			te_eq te1 te2 
		| Type tree_type1, Type tree_type2 ->
			ty_eq  tree_type1 tree_type2
		|_->false)
	
and replace_all tree1 tree2 tree3 =
(******************************
Replaces all instances of tt2 by tt3 in tt1 
tt1[tt3/tt2]
*********************************)
	 
	
	let p_replace_all (left_tree, right_tree) = 
		(replace_all left_tree tree2 tree3, replace_all right_tree tree2 tree3) 
	
	in let replace_ty  = function 
		| T_TyAll (Type var, body) -> 
		(match replace_all body tree2 tree3 with
			|Type tree_body -> make_abstract_type tree_body var
			|_-> raise  TreeFormationError)
			
		| T_TyArr old_pair ->
			(match p_replace_all old_pair with
			|Type t1, Type t2 -> make_arrow_type t1 t2
			|_-> raise TreeFormationError)
		
		| ty_var -> ty_var 
	
	
	in let rec replace_te  te_tree = 
		if tree_eq (Term te_tree) tree2 
		then (match tree3 with Term t3 -> t3 | _-> te_tree)
		else match te_tree with 
		| T_TmAbs (Term var, Term body), _  -> 
			make_abstract_term
				(replace_te body) 
				(match tree2, tree3 with
				| Type _, Type _ -> replace_te var
				|  _ -> var)
	
		| T_TmTAbs (Type var, Term body), _ -> 
			make_abstract_type_term 
				(replace_te body) 
				var
		
		| (T_TmApp (Term left, right)), _ -> 
			make_application_term 
				(replace_te left) 
				(match right with 
					| Term ter_tree_right  -> Term (replace_te ter_tree_right) 
					| Type ty_tree_right ->
						(match tree2, tree3 with
							|Type ty_tree2, Type ty_tree3 -> Type (replace_ty ty_tree_right)
							|_-> right ))
	
		| (T_TmVar (T_BoundedTerm str)), ty_tree -> 
			(T_TmVar (T_BoundedTerm str)), replace_ty ty_tree
		
		| tree_te -> tree_te 	 

		
	in if tree_eq tree1 tree2 then tree3 else
	(match tree1, tree2, tree3 with
		| Type tree_ty1, Type _, Type _ -> Type (replace_ty  tree_ty1)
		| Term tree_te1, Type _, Type _
		| Term tree_te1, Term _, Term _ -> Term (replace_te tree_te1)
		| _ -> tree1)
		
		

and make_abstract_type tree_type1 tree_type2 = 
	let var_node = Type(T_TyVar (T_BoundedType (fresh_ucid( ))))
	in T_TyAll (var_node, 
		replace_all (Type tree_type1) (Type tree_type2) var_node)
	
and make_arrow_type tree_type1 tree_type2 = 
	T_TyArr (Type tree_type1, Type tree_type2)
	
and make_abstract_term (term_tree1, term_type1)  (term_tree2, term_type2) =
(*********************************** 
Builds and abstract term by replacing all occurences 
of te2 in te1 by a bound variable 
************************************)
	let var_node = Term (T_TmVar (T_BoundedTerm (fresh_lcid( ))), term_type2) 
	in T_TmAbs (var_node, replace_all  
			(Term (term_tree1, term_type1))  
			(Term (term_tree2, term_type2))
			var_node), T_TyArr (Type term_type2, Type term_type1)
	
and is_sub_type sub_type_tree sup_type_tree = 
	if tree_eq (Type sub_type_tree) (Type sup_type_tree)
	then false
	else match sup_type_tree with
			| T_TyAll (_, Type ty_tree) -> 
				(is_sub_type sub_type_tree ty_tree) 
			| T_TyArr (Type ty_tree1, Type ty_tree2) -> 
				(is_sub_type sub_type_tree ty_tree1) ||
				(is_sub_type sub_type_tree ty_tree2)
			| _-> false





and type_occurs_in_free_variable (term_tree1, type_tree1) type_tree2 =
	match term_tree1 with
	| T_TmTAbs (_, Term body) 
	| T_TmAbs(_, Term body)->
		type_occurs_in_free_variable body type_tree2
	| T_TmApp(Term left, right)->
		(type_occurs_in_free_variable left type_tree2) || (
			match right with
			|Type _ -> false
			|Term ter_tree_right -> 
				type_occurs_in_free_variable ter_tree_right type_tree2)
	| T_TmVar (T_BoundedTerm _) -> false
	| T_TmVar (T_FreeTerm _)  
	| T_TmVar (T_TmString _)  
	| T_TmVar (T_TmInt _) -> is_sub_type type_tree2 type_tree1
	| _->failwith "type_occurs_in_free_variable "

and make_abstract_type_term (term_tree1, type_tree1)  type_tree2 = 
	if type_occurs_in_free_variable (term_tree1, type_tree1) type_tree2 
		then failwith "In make_abstract_type_term, 
		Type argument occurs in free variable of term"
	
	else let tree_type = make_abstract_type type_tree1 type_tree2 in
	match tree_type with
	| T_TyAll (tyVar, _) ->
			T_TmTAbs( 
				tyVar, 
				replace_all (Term (term_tree1, type_tree1)) (Type type_tree2) tyVar),
			tree_type
	|_-> failwith "In make_abstract_type_term"
	
and apply_type_arg type_tree1 type_tree2 = 
	match type_tree1 with
	|T_TyArr (Type type_tree_left, Type type_tree_right) when
		tree_eq (Type type_tree_left) (Type type_tree2) ->
		type_tree_right
	|T_TyAll (var, body)-> 
		(match replace_all body var (Type type_tree2) with
		|Type type_tree -> type_tree
		|_->failwith "apply_type_arg 1")
	|_->type_tree1

 
and make_application_term (term_tree1, type_tree1)  tree2 = 
	let res_type_tree = 
		apply_type_arg type_tree1 
		(match tree2 with
			| Type tree_ty -> tree_ty
			| Term (_, tree_ty) -> tree_ty) in
	
		T_TmApp (Term (term_tree1, type_tree1) , tree2), res_type_tree
		
		